from plotly.subplots import make_subplots
import plotly.graph_objects as go
import numpy as np


def ColNameSortByNumber(cols):  # 列表通过结尾数值进行排序
    removed_lists = []
    numberList = []
    newList = []
    col_listss = [i for i in cols]

    n = 0
    e_Keys = [
        "sum",
        "Buy",
        "Sell",
        "Main",
        "Boll",
        "day",
        "rolling",
        "by%",
        "SharePrise",
        "cost",
    ]

    for col in col_listss:
        nums = col.split("_")
        num = nums[-1]
        n += 1
        if len(nums) > 1 and num not in e_Keys:
            num = int(num)
            removed_lists.append(col)
            numberList.append(num)
        else:
            newList.append(col)

    newList.sort()
    numberList = sorted(list(set(numberList)))
    for num in numberList:
        for col in removed_lists:
            if int(col.split("_")[-1]) == num:
                newList.append(col)
                # print(num, col)
    return newList


def addTrace(cat, df, colour, fill=None, mode="lines", text=""):
    Temp = go.Scatter(
        x=df["date"],
        y=round(df[cat], 5),
        marker=dict(color=colour, size=3),
        name=cat,
        #   text=df['date'],
        mode=mode,
        #                       # ['none', 'tozeroy', 'tozerox', 'tonexty', 'tonextx', 'toself', 'tonext']
        fill=fill,
        text=text,
        line=dict(
            width=2,
            #  , dash="dash"
        ),
    )
    return Temp


def MainMaTraceAdd(cols, fig, stockdata, Colours, mode="all", Ma_Day_List=[]):
    sum_by, sum_ = 0, 0
    for col in cols:

        if "Price" in col:
            if "in" in col:
                fig.add_trace(
                    addTrace(
                        col,
                        stockdata,
                        "tomato",
                        fill="tozeroy",
                        mode=None,
                        text=stockdata.loc[:, "TAG"],
                    ),
                    row=1,
                    col=1,
                )
            elif "out" in col:
                fig.add_trace(
                    addTrace(
                        col,
                        stockdata,
                        "yellowgreen",
                        fill="tozeroy",
                        mode=None,
                        text=stockdata.loc[:, "TAG"],
                    ),
                    row=1,
                    col=1,
                )
        # elif '-' in col and 'sum' in col:
        #     fig.add_trace(addTrace(col, stockdata, Colours[sum_]),
        #                   row=3, col=1)
        #     sum_ += 1
        # print(Colours[sum_], col)
        elif "_by%" in col and "sum" in col and "close" in col:
            # if "rolling" not in col and "Boll" in col:
            if "Boll" in col:
                # stockdata[col] = stockdata[col] * 0.1
                fig.add_trace(addTrace(col, stockdata, Colours[sum_by]), row=3, col=1)
                sum_by += 1
            # print(Colours[sum_by], col)
        elif "-" in col and "ma" in col:
            if "Main" in col:
                fig.add_trace(addTrace(col, stockdata, Colours[sum_by]), row=3, col=1)
                sum_by += 1
                # print(Colours[Ma_num], col)
            # elif 'Sell' in col:
            #     fig.add_trace(
            #         addTrace(col, stockdata, 'rgb(0,255,50)'),
            #         row=2, col=1)
            # elif 'Buy' in col:
            #     fig.add_trace(
            #         addTrace(col, stockdata, 'rgb(255,100,100)'),
            #         row=2, col=1)

            # elif 'high' in col:
            #     fig.add_trace(addTrace(col, stockdata, colourl[subColor]),
            #                   row=4, col=1)
            # elif 'close' in col:
            #     fig.add_trace(addTrace(col, stockdata, colourl[subColor]),
            #                   row=5, col=1)
        elif "ma" in col and ("Boll" and "by%") not in col:
            if mode != "all":
                if mode in col:
                    stockdata[col].replace(0, np.nan, inplace=True)
                    colNumber = int(col.split("_")[-1])
                    ma_color = Colours[Ma_Day_List.index(colNumber)]
                    fig.add_trace(addTrace(col, stockdata, ma_color), row=1, col=1)
            else:

                colNumber = int(col.split("_")[-1])
                ma_color = Colours[Ma_Day_List.index(colNumber)]
                fig.add_trace(addTrace(col, stockdata, ma_color), row=1, col=1)
        elif "timeing" in col or "day" in col:
            fig.add_trace(addTrace(col, stockdata, Colours[sum_by]), row=3, col=1)
            sum_by += 1
    return fig


def MainRepotDetail(
    report,
    Colours,
    fig,
    colNameLists=["PE_last_year_ratio", "P/E_Price_Earnings_Ratio"],
):
    col = 1
    for cat in colNameLists:
        fig.add_trace(addTrace(cat, report, "gray", mode="lines"), row=5, col=1)
        col += 1
    return fig


def DRAW(df, Report_A, Report_Q, Ma_Day_List=[1], title="", mode="all"):
    """
    绘制
    df = pd.dataframe()
    MA_LIST = Wave.Wave(int)
    keyls = ['close', 'high', 'low', 'open']
    """

    cols = list(set(df.columns.tolist()))
    cols = ColNameSortByNumber(cols)

    row = 3
    tdic = []
    for _ in range(row):
        tdic.append([{"type": "scatter"}])

    # 添加图片框信息
    fig = make_subplots(
        rows=row,
        cols=1,
        shared_xaxes=True,
        vertical_spacing=0.01,
        specs=tdic
        # [
        #     [{"type": "scatter"}],
        #     [{"type": "scatter"}],
        #     [{"type": "scatter"}],
        #     # [{"type": "scatter"}],
        #     #    [{"type": "table"}],
        #     # [{"type": "scatter"}],
        #     #    [{"type": "scatter"}],
        # ]
        ,
        row_heights=[
            0.8,
            0.1,
            0.1,
        ]
        # row_heights=[0.65, 0.1, 0.1, 0.2]
    )

    # plotly 自带的css配色
    CSS_Colours = [
        "yellowgreen",
        "tomato",
        "lightskyblue",
        "lightpink",
        "Aquamarine",
        #    'aliceblue',
        #    'mistyrose',
        "pink",
        "limegreen",
        #    'skyblue',
        "hotpink",
        "cyan",
        "gold",
        "lightsalmon",
        "green",
        "orange",
        "deeppink",
        "Cyan",
        "aqua",
        "aquamarine",
        "azure",
        # 'beige', 'bisque', 'blanchedalmond',
        #    'blueviolet',  'burlywood', 'cadetblue', 'chartreuse', 'chocolate',
        #    'coral', 'cornflowerblue','cornsilk', 'crimson',  'darkblue', 'darkcyan',
        #    'darkgoldenrod', 'darkgray', 'darkgrey', 'darkgreen',
        #    'darkkhaki', 'darkmagenta', 'darkolivegreen', 'darkorange',
        #    'darkorchid', 'darkred', 'darksalmon', 'darkseagreen',
        #    'darkslateblue', 'darkslategray', 'darkslategrey',
        #    'darkturquoise', 'darkviolet',  'deepskyblue',
        #    'dimgray', 'dimgrey', 'dodgerblue', 'firebrick',
        #    'forestgreen', 'fuchsia', 'gainsboro','goldenrod',  'grey',
        #    'greenyellow', 'honeydew',  'indianred', 'indigo',
        #    'ivory', 'khaki', 'lavender', 'lavenderblush', 'lawngreen',
        #    'lemonchiffon', 'lightblue', 'lightcoral', 'lightcyan',
        #    'lightgoldenrodyellow', 'lightgray', 'lightgrey',
        #    'lightgreen', 'lightpink',  'lightseagreen',
        #    'lightslategray', 'lightslategrey',
        #    'lightsteelblue', 'lightyellow', 'lime', 'limegreen',
        #    'linen', 'magenta', 'maroon', 'mediumaquamarine',
        #    'mediumblue', 'mediumorchid', 'mediumpurple',
        #    'mediumseagreen', 'mediumslateblue', 'mediumspringgreen',
        #    'mediumturquoise', 'mediumvioletred', 'midnightblue',
        #    'mintcream',  'moccasin', 'navajowhite', 'navy',
        #    'oldlace', 'orange', 'orangered', 'orchid', 'palegoldenrod', 'palegreen', 'paleturquoise',
        #    'palevioletred', 'papayawhip', 'peachpuff', 'peru',
        #    'plum', 'powderblue', 'purple', 'red', 'rosybrown',
        #    'royalblue', 'rebeccapurple', 'saddlebrown', 'salmon',
        #    'sandybrown', 'seagreen', 'seashell', 'sienna', 'silver',
        #    'slategray', 'slategrey', 'snow','springgreen',  'tan', 'teal', 'thistle',
        #    'turquoise', 'violet', 'wheat','yellow',
    ]
    # 蜡烛图
    fig.add_trace(
        go.Candlestick(
            x=df["date"],
            open=df["open"],
            high=df["high"],
            low=df["low"],
            close=df["close"],
            text=df["date"],
            name="KD",
        ),
        row=1,
        col=1,
    )
    # 绘制MA波动线

    fig = MainMaTraceAdd(cols, fig, df, CSS_Colours, mode=mode, Ma_Day_List=Ma_Day_List)
    # CSS_Colours

    # fig = MainRepotDetail(Report_A, CSS_Colours, fig, colNameLists=[
    #                       'PE_last_year_ratio', 'P/E_Price_Earnings_Ratio'])
    # fig = MainRepotDetail(Report_Q, CSS_Colours, fig, colNameLists=[
    #                       'PE_TTM'])
    try:
        fig.add_trace(
            go.Scatter(
                x=df["date"],
                y=df["MARKER=B"],
                marker=dict(color="tomato", size=15),
                name="MARKER=B",
                mode="markers",
            ),
            row=1,
            col=1,
        )
        fig.add_trace(
            go.Scatter(
                x=df["date"],
                y=df["MARKER=S"],
                marker=dict(color="yellowgreen", size=15),
                name="MARKER=S",
                mode="markers",
            ),
            row=1,
            col=1,
        )

    except Exception as e:
        print(e)

    fig.add_trace(
        addTrace("TotalValues", df, "tomato", mode="lines", fill="tozeroy"),
        row=2,
        col=1,
    )

    fig.add_trace(
        addTrace("AvableCash", df, "green", mode="lines", fill="tozeroy"), row=2, col=1
    )

    fig.add_trace(
        go.Scatter(
            x=df["date"],
            y=round(df["winValue"], 2),
            marker=dict(color="black"),
            name="startCash",
            text=df["newPresent"],
            mode="lines",
            line=dict(width=0.5, dash="dot"),
        ),
        row=2,
        col=1,
    )

    df.loc[:, "zero"] = 0
    for r in range(row):  # 添加0轴
        a = r + 1
        fig.add_trace(
            go.Scatter(
                x=df["date"],
                y=df["zero"],
                marker=dict(color="black"),
                name="zero",
                text=df["date"],
                mode="lines",
                line=dict(width=1, dash="dot"),
            ),
            row=a,
            col=1,
        )

    fig.add_trace(
        go.Scatter(
            x=df["date"],
            y=df["Avg"],
            marker=dict(color="orange"),
            name="Avg",
            mode="lines",
            line=dict(width=5),
        ),
        row=1,
        col=1,
    )

    dtick = (df.loc[:, "close"].max() - df.loc[:, "close"].min()) / 2
    ran = [
        df.loc[:, "close"].min() - dtick * 1.2,
        df.loc[:, "close"].max() + dtick * 1.2,
    ]

    fig.update_layout(
        xaxis_rangeslider_visible=False,
        title=title,
        # paper_bgcolor="rgb(255,255,255)",
        #   xticks=df['date']
        #   template='gridon',xx
        bargroupgap=0,
        yaxis=dict(
            title_text="主图",
            gridwidth=2,
            type="linear",
            range=ran,
            dtick=dtick,
        ),
    )

    fig.update_yaxes(
        title_text="收益情况",
        gridwidth=2,
        dtick=100000,
        row=2,
        col=1,
    )

    fig.update_yaxes(
        dtick=100000,
        range=[-1.05, 1.05],
        gridwidth=2,
        title_text="使用指标",
        row=3,
        col=1,
    )

    df.loc[:, "zero"] = 0
    for r in range(row):  # 添加0轴
        a = r + 1
        fig.update_xaxes(
            range=[df.date.iloc[1], df.date.iloc[-1]],
            showticklabels=False,
            gridwidth=2,
            zeroline=True,
            zerolinewidth=10,
            scaleanchor="x",
            tickmode="array",
            type="category",
            row=a,
            col=1,
        )
    # fig.show()
    title = "_".join(title.split(" "))
    title = "".join(title.split("%"))
    title = "".join(title.split(":"))
    title = "".join(title.split("*"))
    fig.write_html("htmls/{}.html".format(title))
